from pnp.lib.coord import Coord3D
import cv2
import numpy as np
import math

g_allmatches = []
#深度图
g_base_inv = None
#基准特征点 总集
g_base_pts = []
#基准特征点三维坐标 总集
g_base_pts_3d = []
#目标图像特征点
g_target_pts = []

#用于PNP计算
# g_pnp_pt_target=[]
# g_pnp_pt3_base=[]

g_cameraMatrix = None
g_distCoeffs = None
g_residual_threshold = 15
g_residual_threshold_pnp = 10
g_avg_distance= 0

class ModelClass:
    def __init__(self):
        self.rvec = None
        self.tvec = None
        self.good_matches = None
        self.pixel_distances = None
        return


    def estimate(self, src):
        ptspnp3d = []
        ptspnp2d = []
        for idx in src:
            m = g_allmatches[idx][0]
            n = g_allmatches[idx][1]
            if n.distance < m.distance + g_avg_distance:
                return False
            p1 = g_target_pts[m.queryIdx]
            p2 = g_base_pts_3d[m.trainIdx]
            ptspnp3d.append(p2)
            ptspnp2d.append(p1)

        ptspnp3d = np.array(ptspnp3d,dtype=np.float32)
        ptspnp2d = np.array(ptspnp2d,dtype=np.float32)

        retval, rvec, tvec = cv2.solvePnP(ptspnp3d,ptspnp2d,g_cameraMatrix,
                                         g_distCoeffs,flags=cv2.SOLVEPNP_EPNP)
        imgpoints3v, _ = cv2.projectPoints(ptspnp3d,rvec, tvec, g_cameraMatrix, g_distCoeffs)
        imgpoints2v = np.reshape(imgpoints3v,(-1,2))
        error = cv2.norm(ptspnp2d, imgpoints2v, cv2.NORM_L2) / len(imgpoints2v)
        if error > g_residual_threshold_pnp:
            return False

        self.rvec = rvec
        self.tvec = tvec
        self.ptspnp3d = ptspnp3d
        self.ptspnp2d = ptspnp2d
        return True
    #
    # def _match_kernel(self,inindx,inp,points):
    #     mindis = 9999
    #     match = cv2.DMatch()
    #     row = g_match_matrix.shape[0]
    #     col = g_match_matrix.shape[1]
    #     match.queryIdx = -1
    #
    #     for c in range(len(points)):
    #         op = points[c]
    #         dis = math.sqrt(pow(inp[0]-op[0],2)+(pow(inp[1]-op[1],2)))
    #         if dis < 30:
    #             des_dist =  g_match_matrix[inindx,c]
    #             if(des_dist < mindis):
    #                 mindis = des_dist
    #                 match.queryIdx = inindx
    #                 match.trainIdx = c
    #                 match.distance = des_dist
    #
    #     return match

    def _match_kernel2(self,ms,points2v):
        mindis = 9999
        match = None
        outdis = None
        for m in ms:
            r = m.queryIdx
            qpt = g_target_pts[r]
            bpt = points2v[m.trainIdx]
            dis = math.sqrt(pow(qpt[0] - bpt[0], 2) + (pow(qpt[1] - bpt[1], 2)))
            if dis < g_residual_threshold:
                des_dist =  m.distance
                if des_dist < mindis:
                    mindis = des_dist
                    match = m
                    outdis = dis
        return match,outdis

    def residuals(self, src):

        imgpoints3v, _ = cv2.projectPoints(g_base_pts_3d,self.rvec, self.tvec, g_cameraMatrix, g_distCoeffs)
        imgpoints2v = np.reshape(imgpoints3v,(-1,2))
        good_matches = []
        dislist = []
        resarr = np.full(len(g_allmatches),1000,dtype=np.float32)
        ind1 = 0
        for ms in g_allmatches:
            match,dis = self._match_kernel2(ms,imgpoints2v)
            if not match is None:
                resarr[ind1] = match.distance
                good_matches.append(match)
                dislist.append(dis)
            ind1 += 1
        self.good_matches = good_matches
        self.pixel_distances = dislist
        return resarr

        # if len(good_matches) < 1:
        #     return 99999
        # good_matches.sort(key = lambda x:x.distance)
        # alldist = 0
        # num = max(1,min(500,len(good_matches)))
        # for m in range(num):
        #     alldist += good_matches[m].distance
        #
        # return alldist / num

